from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.schema import SystemMessage
from langchain.agents import OpenAIFunctionsAgent
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.agents import AgentExecutor, tool
# from langchain.output_parsers.json import SimpleJsonOutputParser
# from langchain.prompts import PromptTemplate
from flask import Flask, request
import sqlite3
import json


app = Flask(__name__)
SALES_BOT = None
memory = None
tools = []
def init_sales_assistant(vector_store_dir: str = "real_estates_sale"):
    db = FAISS.load_local(vector_store_dir, OpenAIEmbeddings(), allow_dangerous_deserialization=True)
    llm = ChatOpenAI(model_name = 'gpt-3.5-turbo', temperature = 0)
    retriever = db.as_retriever(search_type="similarity_score_threshold",
                    search_kwargs={"score_threshold": 0.8})

    global SALES_BOT
    # 创建memory
    global memory
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key="answer",
        return_messages=True
    )

    # 创建ConversationalRetrievalChain
    SALES_BOT = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory
    )

    SALES_BOT.return_source_documents = True

    return SALES_BOT


@app.route('/api/chat', methods=['POST', 'GET'])
def sales_chat():
    messages = request.get_json()

    message = messages.get('message')
    history = messages.get("history")
    print(f"Message: {message}")
    print(f"History: {history}")

    enable_chat = True

    answers = SALES_BOT({"question": message})
    print(answers)
    if (answers["source_documents"]) or enable_chat:
        print(f"[result]{answers['answer']}")
        print(f"[source_documents]: {answers['source_documents']}")
        return answers["answer"]
    else:
        return "对不起，我暂时无法回答的您的问题！"


@tool
def query_product(brand: str = '', category: str = '', feature: str = ''):
    """返回查询到的商品."""
    print('解构的查询条件：', brand, category, feature)
    conn = sqlite3.connect('./data/sales.db')
    cursor = conn.cursor()
    sql = "select * from home_appliances where 1=1 "
    if brand != '':
        sql = sql + " and brand like '%" + brand + "%'"
    if category != '':
        sql = sql + " and category like '%" + category + "%'"
    if feature != '':
        sql = sql + " and feature like '%" + feature + "%'"
    sql = sql + ' limit 3 offset 0'
    print(sql)
    cursor.execute(sql)
    # 获取查询结果的列名
    columns = [description[0] for description in cursor.description]
    # 获取查询结果的所有行，并转换成字典列表
    rows = [dict(zip(columns, row)) for row in cursor.fetchall()]
    print(rows)
    # 将字典列表转换成JSON字符串
    json_data = json.dumps(rows, indent=4)  # indent=4 使得输出更易读
    print(json_data)
    return json_data


# 推荐商品接口
@app.route('/api/recommendation', methods=['POST', 'GET'])
def recommendation():
    tools = [query_product]
    history_messages = request.get_json()
    chat_memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    build_memory(history_messages, chat_memory)
    llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0)
    system_message = SystemMessage(content="""
        你是一个家电销售助理，根据当前对话记录，推断是否需要给客户推荐商品。若需要请从对话中提取出来接下来要给客户推荐商品的查询条件，并查询推荐的商品。
        支持的条件有：
            brand: 家电品牌,如格力、美的等
            category: 家电种类，如洗衣机、空调、电视等
            feature: 家电特性，如静音、智能、美观等
        请以markdown格式回答用户的提问    
        """)
    prompt = OpenAIFunctionsAgent.create_prompt(system_message=system_message)
    agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, memory=chat_memory, verbose=True)

    # output_parser = SimpleJsonOutputParser()
    # template = """
    # {question}
    #
    # {format_instructions}"""
    #
    # prompt = PromptTemplate.from_template(
    #     template,
    #     partial_variables={"format_instructions": output_parser.get_format_instructions()},
    # )
    # _input = prompt.format(question='以上是对话的记录，请根据最后一次的对话记录确认是否需要给客户推荐商品，若需要请查询推荐的商品', verbose = True)
    result = agent_executor.run('以上是对话的记录，请根据最后一次的对话记录确认是否需要给客户推荐商品，若需要请查询推荐的商品')
    print(result)
    return result


def build_memory(history_messages, chat_memory):
    messages = history_messages['message']
    i = 0
    while i < len(messages):
        chat_memory.save_context({ "input": messages[i]['message'] },
                                 { "output": messages[i+1]['message'] if i + 1 < len(messages) else ''})
        i = i + 2


if __name__ == "__main__":
    init_sales_assistant()
    app.run(host='0.0.0.0', port=8888, debug=True)
